Description

Use multi-tensor kernel strategy for resetting gradients.
Undo #16716

Checklist

Essentials

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage (https://github.com/sxjscience/mxnet/blob/master/tests/python/unittest/test_numpy_gluon.py#L185)
• Code is well-documented

Changes

• in zero_grad method, use ndarray.reset_arrays for resetting multiple arrays within same kernel

Comments

Performance observed training BERT-large on single V100 GPU:

........................................................................Throughput samples / s
BatchSize - BatchAccumulation............Pre-Change............This-PR (multi-tensor )............ Improvement(%)
................4 - 2 .............................................. 29.82 ..........................36.24.............................................. 21.5
................4 - 4............................................... 36.89............................41.45..............................................12.3
................8 - 4............................................... 45.25............................48.41.............................................. 6.9